import numpy.matlib
import numpy as np
import pandas as pd
import statsmodels.api as sm

from .returns import Returns


class FactorModel:
    def __init__(self) -> None:
        # start & end date
        self.start_date = ""
        self.end_date = ""
        # risk free rate
        self.risk_free_rate = None
        self.rf_set = False
        # factor matrix
        self.factor_matrix = pd.DataFrame()
        # all factors
        self.factor_list = []
        # all factors (use as backup of factors)
        self.factors = {}
        # cov matrix
        self.covmatrix = None
        pass

    def set_risk_free_rate(self, rf: Returns) -> None:
        """Set risk free rate of the factor model
        Risk free rate is a special FactorReturn.

        Args:
            rf (FactorReturn): the risk free rate.
        """
        rf.name = "_rf"
        self.risk_free_rate = rf
        self.add_factor_return(rf, False)
        self.rf_set = True

    def add_factor_return(
        self, factor_return: Returns, adjust_rf: bool = False
    ) -> None:
        factor_name = factor_return.name
        if not factor_name.startswith("_"):
            self.factor_list.append(factor_name)
        self.factors[factor_name] = factor_return
        data_slice = factor_return[
            (factor_return.index >= self.start_date)
            & (factor_return.index <= self.end_date)
        ]
        join_type = "inner" if len(self.factor_matrix.columns) else "outer"
        self.factor_matrix = pd.concat(
            [self.factor_matrix, data_slice], join=join_type, axis=1
        )
        if adjust_rf:
            if not self.rf_set:
                print("Must set risk free rate")
                return None
            self.factor_matrix[factor_name] -= self.factor_matrix["_rf"]
        return None

    def calculate_cov_matrix(self) -> None:
        factors = len(self.factor_list)
        self.covmatrix = np.matlib.zeros((factors, factors))
        for i in range(factors):
            for j in range(factors):
                i_name = self.factor_list[i]
                j_name = self.factor_list[j]
                cov = np.cov(self.factor_matrix[i_name], self.factor_matrix[j_name])
                self.covmatrix[i, j] = cov[0, 1]
        return None

    def get_asset_exposure(
        self, asset_return: Returns, alpha: bool = False, adjust_rf: bool = True
    ) -> dict:
        xy = pd.concat([self.factor_matrix, asset_return], join="inner", axis=1)

        if alpha:
            x = sm.add_constant(xy[self.factor_list])
        else:
            x = xy[self.factor_list]

        y = xy[asset_return.name]
        if adjust_rf:
            y -= xy["_rf"]

        model = sm.OLS(y, x)
        res = model.fit()
        param = res.params.to_dict()
        exposure = {
            "exposure": {x: param[x] for x in self.factor_list},
            "indivial_risk": np.sqrt(res.mse_resid),
        }
        if alpha:
            exposure["alpha"] = param["const"]
        else:
            exposure["alpha"] = 0
        return exposure
